import pandas as pd
import subprocess
import json
import re
import time

# ==============================
# ✅ 配置区
# ==============================
MODEL_NAME = "qwen2"
INPUT_FILE = r"D:\software\Desktop\Ollama_damo\test_extract.xlsx"
OUTPUT_FILE = "fine_grained_triples_chinese2.csv"

# 中文 Schema
SCHEMA = {
    "地理位置": "地理位置（区/县）",
    "经纬度坐标": "经纬度坐标",
    "开放时间": "开放时间",
    "闭馆日": "闭馆日",
    "免票/优惠政策": "免票/优惠政策",
    "无障碍设施情况": "无障碍设施情况",
    "周边景点": "周边景点",
    "建议游览时长": "建议游览时长",
    "拍照/打卡点": "拍照/打卡点",
    "所属历史时期": "所属历史时期",
    "是否为世界文化遗产": "是否为世界文化遗产",
    "周边美食": "周边美食",
    "是否适合夜游": "是否适合夜游",
    "适合人群": "适合人群"
}

# ==============================
# ✅ 工具函数
# ==============================
def call_qwen2(prompt, timeout=45):
    """调用本地 Ollama Qwen2"""
    try:
        result = subprocess.run(
            ["ollama", "run", MODEL_NAME],
            input=prompt,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
            encoding='utf-8'
        )
        return result.stdout.strip()
    except subprocess.TimeoutExpired:
        print("⚠️ 调用超时，跳过此项")
        return ""
    except Exception as e:
        print(f"[ERROR] {e}")
        return ""

def safe_json_parse(text):
    """尝试清洗并解析模型返回的JSON"""
    try:
        clean = re.sub(r"```json|```", "", text).strip()
        clean = re.sub(r"```.*?```", "", clean, flags=re.DOTALL)
        return json.loads(clean)
    except Exception as e:
        print(f"[WARN] JSON解析失败: {e}\n原始输出: {text[:200]}")
        return []

# ==============================
# ✅ 模型调用函数
# ==============================
def enrich_attributes(subject, context_text):
    """根据中文 schema 扩充属性"""
    prompt = f"""
你是一名陕西文旅知识图谱专家。请根据以下信息补全景点“{subject}”的结构化属性。

信息文本：
{context_text}

请按照以下中文 schema 提取：
{json.dumps(SCHEMA, ensure_ascii=False, indent=2)}

输出JSON列表，格式示例：
[
  {{"predicate": "地理位置", "object": "西安市临潼区"}},
  {{"predicate": "开放时间", "object": "08:30-17:00"}}
]
若无信息则返回 []。
"""
    resp = call_qwen2(prompt)
    data = safe_json_parse(resp)
    return [(subject, d["predicate"], d["object"]) for d in data if "predicate" in d and "object" in d]

def split_ticket_options(subject, ticket_text):
    """拆分票务信息"""
    prompt = f"""
你是一名陕西文旅知识图谱专家。
请将景点“{subject}”的票务信息转为结构化JSON，示例：
[
  {{
    "票种": "成人票",
    "价格": 120,
    "币种": "CNY"
  }},
  {{
    "票种": "学生票",
    "价格": 60,
    "币种": "CNY",
    "凭证": "学生证"
  }}
]
信息如下：
{ticket_text}
"""
    resp = call_qwen2(prompt)
    data = safe_json_parse(resp)
    triples = []
    for t in data:
        triples.append((subject, "票务选项", json.dumps(t, ensure_ascii=False)))
    return triples

def split_traffic_info(subject, traffic_text):
    """结构化交通信息"""
    prompt = f"""
请提取景点“{subject}”的交通出行方式，并转为结构化JSON：
[
  {{"方式": "地铁", "线路": "1号线", "站点": "临潼站"}},
  {{"方式": "公交", "路线": "306路", "备注": "兵马俑专线"}}
]
信息如下：
{traffic_text}
"""
    resp = call_qwen2(prompt)
    data = safe_json_parse(resp)
    triples = []
    for t in data:
        triples.append((subject, "交通方式", json.dumps(t, ensure_ascii=False)))
    return triples

# ==============================
# ✅ 主流程
# ==============================
def main():
    print("📂 正在加载文件...")
    df = pd.read_excel(INPUT_FILE)
    print(f"✅ 文件加载成功，共 {len(df)} 条数据。开始调用模型...")

    all_triples = []

    for idx, row in df.iterrows():
        subj = str(row["subject"])
        rel = str(row["predicate"])
        obj = str(row["object"])

        # 针对不同字段执行不同策略
        if rel.lower() == "ticket":
            new_triples = split_ticket_options(subj, obj)
        elif rel.lower() == "traffic":
            new_triples = split_traffic_info(subj, obj)
        else:
            new_triples = enrich_attributes(subj, f"{rel}: {obj}")

        all_triples.extend(new_triples)
        print(f"[{idx+1}/{len(df)}] ✅ {subj} -> {len(new_triples)} 条新属性")
        time.sleep(1)  # 控制调用频率，防止模型阻塞

    # 保存结果
    result_df = pd.DataFrame(all_triples, columns=["subject", "predicate", "object"])
    result_df.to_csv(OUTPUT_FILE, index=False, encoding="utf-8-sig")
    print(f"\n🎯 处理完成，共生成 {len(result_df)} 条细粒度三元组，已保存到：{OUTPUT_FILE}")

# ==============================
# ✅ 启动入口
# ==============================
if __name__ == "__main__":
    main()
